# 导入必要的库
from langchain_core.messages import merge_content
from langgraph.graph import StateGraph, END
from typing import Dict, Any, TypedDict

from rag.hello_rag import retriever


# 定义状态类型 - 使用TypedDict来明确状态结构
class GraphState(TypedDict):
    name: str
    question: str
    greeting: str
    answer: str
    farewell: str


# 定义第一个节点函数 - 生成问候语
def generate_greeting(state: GraphState) -> GraphState:
    """生成问候语的节点"""
    name = state.get("name", "访客")
    greeting = f"你好，{name}！欢迎使用LangGraph。"
    print(f"生成问候: {greeting}")
    return {**state, "greeting": greeting}


# 定义第二个节点函数 - 处理问题
def process_question(state: GraphState) -> GraphState:
    """处理问题并生成回答的节点"""
    question = state.get("question", "")
    answer = f"关于'{question}'，这是一个很好的问题！LangGraph是用于构建AI工作流的工具。"
    print(f"处理问题: {answer}")
    return {**state, "answer": answer}


# 定义第三个节点函数 - 生成结束语
def generate_farewell(state: GraphState) -> GraphState:
    """生成结束语的节点"""
    farewell = "感谢使用LangGraph，祝你学习愉快！"
    print(f"生成结束语: {farewell}")
    return {**state, "farewell": farewell}


# 创建图的工作流
def create_workflow():
    """创建并配置工作流图"""

    # 初始化图 - 使用StateGraph而不是Graph
    workflow = StateGraph(GraphState)

    # 添加节点到图中
    workflow.add_node("greet", generate_greeting)  # 问候节点
    workflow.add_node("process", process_question)  # 处理问题节点
    workflow.add_node("farewell", generate_farewell)  # 结束节点

    # 设置入口点 - 从问候开始
    workflow.set_entry_point("greet")

    # 定义节点间的流转关系
    workflow.add_edge("greet", "process")  # 问候完成后转到问题处理
    workflow.add_edge("process", "farewell")  # 问题处理完成后转到结束
    workflow.add_edge("farewell", END)  # 结束节点后流程终止

    # 编译图
    return workflow.compile()


# 主函数
def main():
    """主执行函数"""

    # 创建并编译工作流
    app = create_workflow()

    # 定义初始状态
    initial_state = {
        "name": "小明",
        "question": "什么是LangGraph？",
        "greeting": "",
        "answer": "",
        "farewell": ""
    }

    # 执行工作流
    print("=== LangGraph 工作流执行开始 ===")
    result = app.invoke(initial_state)

    # 打印最终结果
    print("\n=== 最终执行结果 ===")
    for key, value in result.items():
        print(f"{key}: {value}")


# 程序入口
if __name__ == "__main__":
    main()
